# coding=utf-8
# Adapted from https://github.com/HengyiWang/spann3r/blob/main/spann3r/tools/eval_recon.py
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# Copyright (C) HengyiWang. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
import numpy as np
from scipy.spatial import cKDTree as KDTree


def accuracy(gt_points, rec_points, gt_normals=None, rec_normals=None):
    gt_points_kd_tree = KDTree(gt_points)
    distances, idx = gt_points_kd_tree.query(rec_points, workers=-1)
    acc = np.mean(distances)

    acc_median = np.median(distances)

    if gt_normals is not None and rec_normals is not None:
        normal_dot = np.sum(gt_normals[idx] * rec_normals, axis=-1)
        normal_dot = np.abs(normal_dot)

        return acc, acc_median, np.mean(normal_dot), np.median(normal_dot)

    return acc, acc_median


def completion(gt_points, rec_points, gt_normals=None, rec_normals=None):
    gt_points_kd_tree = KDTree(rec_points)
    distances, idx = gt_points_kd_tree.query(gt_points, workers=-1)
    comp = np.mean(distances)
    comp_median = np.median(distances)

    if gt_normals is not None and rec_normals is not None:
        normal_dot = np.sum(gt_normals * rec_normals[idx], axis=-1)
        normal_dot = np.abs(normal_dot)

        return comp, comp_median, np.mean(normal_dot), np.median(normal_dot)

    return comp, comp_median


def calc_performance(pcd, pcd_gt):
    pcd.estimate_normals()
    pcd_gt.estimate_normals()

    gt_normal = np.asarray(pcd_gt.normals)
    pred_normal = np.asarray(pcd.normals)

    acc, _, nc1, _ = accuracy(
        pcd_gt.points, pcd.points, gt_normal, pred_normal
    )
    comp, _, nc2, _ = completion(
        pcd_gt.points, pcd.points, gt_normal, pred_normal
    )

    return acc, nc1, comp, nc2